import numpy as np
from keras._tf_keras.keras.datasets import mnist
from keras.api.utils import to_categorical
from keras import models
from keras import layers

def load_data():
    (x_train, y_train), (x_test, y_test) = mnist.load_data()
    print("x_train.shape ", x_train.shape)
    print("y_train.shape", y_train.shape)
    print("x_test.shape ", x_test.shape)
    print("y_test.shape ", y_test.shape)

    # reshape
    x_train = x_train.reshape([60000,784])
    x_test = x_test.reshape([10000,784])
    print("x_train.shape ", x_train.shape)
    print("x_test.shape ", x_test.shape)
    x_train = x_train.astype('float32')
    x_test = x_test.astype('float32')
    print("x_train.type ", x_train.dtype)
    print("x_test.type ", x_test.dtype)

    # 归一化
    x_train /= 255
    x_test /= 255

    print(x_train[1][100:150])

    # one-hot encode
    train_labels = to_categorical(y_train)
    test_labels = to_categorical(y_test)
    print("train_label[0:4] ",train_labels[:5])

    return x_train, train_labels, x_test, test_labels


def train(x_train, train_labels, x_test, test_labels, epochs=10, batch_size=128, validation_split=0.2):
    # build model
    model = models.Sequential()
    model.add(layers.Dense(128, input_shape=(28*28,), activation='relu'))
    model.add(layers.Dense(128, activation='relu'))
    model.add(layers.Dense(10, activation='softmax'))
    print(model.summary())
    # 编译模型
    model.compile(loss='categorical_crossentropy',optimizer='sgd',metrics=['accuracy'])
    # 训练模型
    model.fit(x_train, train_labels, epochs=epochs, batch_size=batch_size, validation_split=validation_split, verbose=1)
    test_loss, test_acc = model.evaluate(x_test, test_labels, verbose=1)

    print('Test loss:', test_loss)
    print('Test accuracy:', test_acc)

    return model


if __name__ == '__main__':
    x_train, train_labels, x_test, test_labels = load_data()
    model = train(x_train, train_labels, x_test, test_labels, epochs=20, batch_size=128, validation_split=0.2)
